#!/usr/bin/env python3

#==============================================================================
# Functionality
#==============================================================================
import pdb
import sys
import os
import re

# utility funcs, classes, etc go here.

def asserting(cond):
    if not cond:
        pdb.set_trace()
    assert(cond)

def has_stdin():
    return not sys.stdin.isatty()

def reg(pat, flags=0):
    return re.compile(pat, re.VERBOSE | flags)

#==============================================================================
# Cmdline
#==============================================================================
import argparse

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, 
    description="""
TODO
""")
     
parser.add_argument('-v', '--verbose',
    action="store_true",
    help="verbose output" )
     
parser.add_argument('-d', '--delay',
    type=float,
    default=2.0,
    help="time to wait after using webhook" )
     
parser.add_argument('-w', '--webhook',
    type=str,
    default="",
    help="The discord webhook URL" )
     
parser.add_argument('-c', '--compress',
    action="store_true",
    help="Upload as JPEG with quality=95" )

args = None

#==============================================================================
# Main
#==============================================================================

from base64 import b64decode
import re
#from tensorboard.compat.proto import event_pb2
#import google.protobuf.json_format as jf
import json
import hashlib
import subprocess
import requests
import functools
from io import StringIO
from io import BytesIO
import time
import PIL.Image

db = {}
regex = reg(r""" "encodedImageString" \s* : \s* "([^"]+?)" """)

def replace_upload(m, *, line=None):
  #import pdb; pdb.set_trace()
  img = b64decode(m.group(1))
  h = hashlib.new('md5')
  h.update(img)
  md = h.hexdigest()
  filename = md + ('.jpg' if args.compress else '.png')
  if args.compress:
    # convert to jpg
    image = PIL.Image.open(BytesIO(img))
    bio = BytesIO()
    image.save(bio, format='jpeg', quality=95)
    img = bio.getvalue()
  #db[md] = img
  req = requests.get('https://test.tensorfork.com/webdis/GET/{}.txt'.format(md))
  if req.status_code != 200:
    processed = re.sub(regex, '"md5sum": "{}"'.format(md), line)
    try:
      out = subprocess.run(['discord-webhook-upload-file', args.webhook, '', filename, json.dumps("```{}```".format(processed))[1:-1]], input=img, check=True, stdout=subprocess.PIPE)
      zz = out.stdout.decode('utf8')
      js = zz.splitlines()[-1]
      files = {'upload_file': StringIO(js)}
      req = requests.put('https://test.tensorfork.com/webdis/SET/{}.txt'.format(md), data=BytesIO(js.encode('utf8')))
      time.sleep(args.delay)
    except:
      import traceback
      traceback.print_exc()
      return '"error": "%d"' % req.status_code
  else:
    js = req.text
  j = json.loads(js)
  url = j['attachments'][0]['url']
  #import pdb; pdb.set_trace()

  #return '"md5sum": "{}"'.format(md)
  return '"md5sum": "{}", "url": "{}"'.format(md, url)
  

def process_json_events(f):
  for line in f:
    line = re.sub(regex, functools.partial(replace_upload, line=line), line)
    sys.stdout.write(line)
    sys.stdout.flush()


def run():
    if args.verbose:
        print(args)
    if len(args.args) <= 0 and not has_stdin():
        # if there were no args and there was no input, prompt user.
        print('Enter input (press Ctrl-D when done):')
    if len(args.args) <= 0 or has_stdin():
        process_json_events(sys.stdin)
        return
    else:
        # for each arg on cmdline...
        for arg in args.args:
          with open(arg) as f:
            process_json_events(f)

def main():
    try:
        global args
        if not args:
            args, leftovers = parser.parse_known_args()
            args.args = leftovers
        return run()
    except IOError:
        # http://stackoverflow.com/questions/15793886/how-to-avoid-a-broken-pipe-error-when-printing-a-large-amount-of-formatted-data
        try:
            sys.stdout.close()
        except IOError:
            pass
        try:
            sys.stderr.close()
        except IOError:
            pass

if __name__ == "__main__":
    main()

